{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d159d5b6-d456-46ba-8eb6-b88217de2940",
   "metadata": {},
   "source": [
    "# Vanna AI LlamaPack\n",
    "\n",
    "Vanna AI is an open-source RAG framework for SQL generation. It works in two steps:\n",
    "1. Train a RAG model on your data\n",
    "2. Ask questions (use reference corpus to generate SQL queries that can run on your db).\n",
    "\n",
    "Check out the [Github project](https://github.com/vanna-ai/vanna) and the [docs](https://vanna.ai/docs/) for more details.\n",
    "\n",
    "This LlamaPack creates a simple `VannaQueryEngine` with vanna, ChromaDB and OpenAI, and allows you to train and ask questions over a SQL database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "24495652-d1e8-48a6-916f-0c947605c5b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Option: if developing with the llama_hub package\n",
    "# from llama_hub.llama_packs.vanna.base import VannaPack\n",
    "\n",
    "# Option: download_llama_pack\n",
    "from llama_index.llama_pack import download_llama_pack\n",
    "\n",
    "VannaPack = download_llama_pack(\n",
    "    \"VannaPack\",\n",
    "    \"./vanna_pack\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e544a40d-1283-4778-aa44-eb4cb801b506",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Connected to db: chinook.db\n",
      "> Training on CREATE TABLE \"albums\"\n",
      "(\n",
      "    [AlbumId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [Title] NVARCHAR(160)  NOT NULL,\n",
      "    [ArtistId] INTEGER  NOT NULL,\n",
      "    FOREIGN KEY ([ArtistId]) REFERENCES \"artists\" ([ArtistId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "Adding ddl: CREATE TABLE \"albums\"\n",
      "(\n",
      "    [AlbumId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [Title] NVARCHAR(160)  NOT NULL,\n",
      "    [ArtistId] INTEGER  NOT NULL,\n",
      "    FOREIGN KEY ([ArtistId]) REFERENCES \"artists\" ([ArtistId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "> Training on CREATE TABLE sqlite_sequence(name,seq)\n",
      "Adding ddl: CREATE TABLE sqlite_sequence(name,seq)\n",
      "> Training on CREATE TABLE \"artists\"\n",
      "(\n",
      "    [ArtistId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [Name] NVARCHAR(120)\n",
      ")\n",
      "Adding ddl: CREATE TABLE \"artists\"\n",
      "(\n",
      "    [ArtistId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [Name] NVARCHAR(120)\n",
      ")\n",
      "> Training on CREATE TABLE \"customers\"\n",
      "(\n",
      "    [CustomerId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [FirstName] NVARCHAR(40)  NOT NULL,\n",
      "    [LastName] NVARCHAR(20)  NOT NULL,\n",
      "    [Company] NVARCHAR(80),\n",
      "    [Address] NVARCHAR(70),\n",
      "    [City] NVARCHAR(40),\n",
      "    [State] NVARCHAR(40),\n",
      "    [Country] NVARCHAR(40),\n",
      "    [PostalCode] NVARCHAR(10),\n",
      "    [Phone] NVARCHAR(24),\n",
      "    [Fax] NVARCHAR(24),\n",
      "    [Email] NVARCHAR(60)  NOT NULL,\n",
      "    [SupportRepId] INTEGER,\n",
      "    FOREIGN KEY ([SupportRepId]) REFERENCES \"employees\" ([EmployeeId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "Adding ddl: CREATE TABLE \"customers\"\n",
      "(\n",
      "    [CustomerId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [FirstName] NVARCHAR(40)  NOT NULL,\n",
      "    [LastName] NVARCHAR(20)  NOT NULL,\n",
      "    [Company] NVARCHAR(80),\n",
      "    [Address] NVARCHAR(70),\n",
      "    [City] NVARCHAR(40),\n",
      "    [State] NVARCHAR(40),\n",
      "    [Country] NVARCHAR(40),\n",
      "    [PostalCode] NVARCHAR(10),\n",
      "    [Phone] NVARCHAR(24),\n",
      "    [Fax] NVARCHAR(24),\n",
      "    [Email] NVARCHAR(60)  NOT NULL,\n",
      "    [SupportRepId] INTEGER,\n",
      "    FOREIGN KEY ([SupportRepId]) REFERENCES \"employees\" ([EmployeeId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "> Training on CREATE TABLE \"employees\"\n",
      "(\n",
      "    [EmployeeId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [LastName] NVARCHAR(20)  NOT NULL,\n",
      "    [FirstName] NVARCHAR(20)  NOT NULL,\n",
      "    [Title] NVARCHAR(30),\n",
      "    [ReportsTo] INTEGER,\n",
      "    [BirthDate] DATETIME,\n",
      "    [HireDate] DATETIME,\n",
      "    [Address] NVARCHAR(70),\n",
      "    [City] NVARCHAR(40),\n",
      "    [State] NVARCHAR(40),\n",
      "    [Country] NVARCHAR(40),\n",
      "    [PostalCode] NVARCHAR(10),\n",
      "    [Phone] NVARCHAR(24),\n",
      "    [Fax] NVARCHAR(24),\n",
      "    [Email] NVARCHAR(60),\n",
      "    FOREIGN KEY ([ReportsTo]) REFERENCES \"employees\" ([EmployeeId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "Adding ddl: CREATE TABLE \"employees\"\n",
      "(\n",
      "    [EmployeeId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [LastName] NVARCHAR(20)  NOT NULL,\n",
      "    [FirstName] NVARCHAR(20)  NOT NULL,\n",
      "    [Title] NVARCHAR(30),\n",
      "    [ReportsTo] INTEGER,\n",
      "    [BirthDate] DATETIME,\n",
      "    [HireDate] DATETIME,\n",
      "    [Address] NVARCHAR(70),\n",
      "    [City] NVARCHAR(40),\n",
      "    [State] NVARCHAR(40),\n",
      "    [Country] NVARCHAR(40),\n",
      "    [PostalCode] NVARCHAR(10),\n",
      "    [Phone] NVARCHAR(24),\n",
      "    [Fax] NVARCHAR(24),\n",
      "    [Email] NVARCHAR(60),\n",
      "    FOREIGN KEY ([ReportsTo]) REFERENCES \"employees\" ([EmployeeId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "> Training on CREATE TABLE \"genres\"\n",
      "(\n",
      "    [GenreId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [Name] NVARCHAR(120)\n",
      ")\n",
      "Adding ddl: CREATE TABLE \"genres\"\n",
      "(\n",
      "    [GenreId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [Name] NVARCHAR(120)\n",
      ")\n",
      "> Training on CREATE TABLE \"invoices\"\n",
      "(\n",
      "    [InvoiceId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [CustomerId] INTEGER  NOT NULL,\n",
      "    [InvoiceDate] DATETIME  NOT NULL,\n",
      "    [BillingAddress] NVARCHAR(70),\n",
      "    [BillingCity] NVARCHAR(40),\n",
      "    [BillingState] NVARCHAR(40),\n",
      "    [BillingCountry] NVARCHAR(40),\n",
      "    [BillingPostalCode] NVARCHAR(10),\n",
      "    [Total] NUMERIC(10,2)  NOT NULL,\n",
      "    FOREIGN KEY ([CustomerId]) REFERENCES \"customers\" ([CustomerId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "Adding ddl: CREATE TABLE \"invoices\"\n",
      "(\n",
      "    [InvoiceId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [CustomerId] INTEGER  NOT NULL,\n",
      "    [InvoiceDate] DATETIME  NOT NULL,\n",
      "    [BillingAddress] NVARCHAR(70),\n",
      "    [BillingCity] NVARCHAR(40),\n",
      "    [BillingState] NVARCHAR(40),\n",
      "    [BillingCountry] NVARCHAR(40),\n",
      "    [BillingPostalCode] NVARCHAR(10),\n",
      "    [Total] NUMERIC(10,2)  NOT NULL,\n",
      "    FOREIGN KEY ([CustomerId]) REFERENCES \"customers\" ([CustomerId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "> Training on CREATE TABLE \"invoice_items\"\n",
      "(\n",
      "    [InvoiceLineId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [InvoiceId] INTEGER  NOT NULL,\n",
      "    [TrackId] INTEGER  NOT NULL,\n",
      "    [UnitPrice] NUMERIC(10,2)  NOT NULL,\n",
      "    [Quantity] INTEGER  NOT NULL,\n",
      "    FOREIGN KEY ([InvoiceId]) REFERENCES \"invoices\" ([InvoiceId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION,\n",
      "    FOREIGN KEY ([TrackId]) REFERENCES \"tracks\" ([TrackId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "Adding ddl: CREATE TABLE \"invoice_items\"\n",
      "(\n",
      "    [InvoiceLineId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [InvoiceId] INTEGER  NOT NULL,\n",
      "    [TrackId] INTEGER  NOT NULL,\n",
      "    [UnitPrice] NUMERIC(10,2)  NOT NULL,\n",
      "    [Quantity] INTEGER  NOT NULL,\n",
      "    FOREIGN KEY ([InvoiceId]) REFERENCES \"invoices\" ([InvoiceId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION,\n",
      "    FOREIGN KEY ([TrackId]) REFERENCES \"tracks\" ([TrackId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "> Training on CREATE TABLE \"media_types\"\n",
      "(\n",
      "    [MediaTypeId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [Name] NVARCHAR(120)\n",
      ")\n",
      "Adding ddl: CREATE TABLE \"media_types\"\n",
      "(\n",
      "    [MediaTypeId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [Name] NVARCHAR(120)\n",
      ")\n",
      "> Training on CREATE TABLE \"playlists\"\n",
      "(\n",
      "    [PlaylistId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [Name] NVARCHAR(120)\n",
      ")\n",
      "Adding ddl: CREATE TABLE \"playlists\"\n",
      "(\n",
      "    [PlaylistId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [Name] NVARCHAR(120)\n",
      ")\n",
      "> Training on CREATE TABLE \"playlist_track\"\n",
      "(\n",
      "    [PlaylistId] INTEGER  NOT NULL,\n",
      "    [TrackId] INTEGER  NOT NULL,\n",
      "    CONSTRAINT [PK_PlaylistTrack] PRIMARY KEY  ([PlaylistId], [TrackId]),\n",
      "    FOREIGN KEY ([PlaylistId]) REFERENCES \"playlists\" ([PlaylistId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION,\n",
      "    FOREIGN KEY ([TrackId]) REFERENCES \"tracks\" ([TrackId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "Adding ddl: CREATE TABLE \"playlist_track\"\n",
      "(\n",
      "    [PlaylistId] INTEGER  NOT NULL,\n",
      "    [TrackId] INTEGER  NOT NULL,\n",
      "    CONSTRAINT [PK_PlaylistTrack] PRIMARY KEY  ([PlaylistId], [TrackId]),\n",
      "    FOREIGN KEY ([PlaylistId]) REFERENCES \"playlists\" ([PlaylistId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION,\n",
      "    FOREIGN KEY ([TrackId]) REFERENCES \"tracks\" ([TrackId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "> Training on CREATE TABLE \"tracks\"\n",
      "(\n",
      "    [TrackId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [Name] NVARCHAR(200)  NOT NULL,\n",
      "    [AlbumId] INTEGER,\n",
      "    [MediaTypeId] INTEGER  NOT NULL,\n",
      "    [GenreId] INTEGER,\n",
      "    [Composer] NVARCHAR(220),\n",
      "    [Milliseconds] INTEGER  NOT NULL,\n",
      "    [Bytes] INTEGER,\n",
      "    [UnitPrice] NUMERIC(10,2)  NOT NULL,\n",
      "    FOREIGN KEY ([AlbumId]) REFERENCES \"albums\" ([AlbumId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION,\n",
      "    FOREIGN KEY ([GenreId]) REFERENCES \"genres\" ([GenreId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION,\n",
      "    FOREIGN KEY ([MediaTypeId]) REFERENCES \"media_types\" ([MediaTypeId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "Adding ddl: CREATE TABLE \"tracks\"\n",
      "(\n",
      "    [TrackId] INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n",
      "    [Name] NVARCHAR(200)  NOT NULL,\n",
      "    [AlbumId] INTEGER,\n",
      "    [MediaTypeId] INTEGER  NOT NULL,\n",
      "    [GenreId] INTEGER,\n",
      "    [Composer] NVARCHAR(220),\n",
      "    [Milliseconds] INTEGER  NOT NULL,\n",
      "    [Bytes] INTEGER,\n",
      "    [UnitPrice] NUMERIC(10,2)  NOT NULL,\n",
      "    FOREIGN KEY ([AlbumId]) REFERENCES \"albums\" ([AlbumId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION,\n",
      "    FOREIGN KEY ([GenreId]) REFERENCES \"genres\" ([GenreId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION,\n",
      "    FOREIGN KEY ([MediaTypeId]) REFERENCES \"media_types\" ([MediaTypeId]) \n",
      "\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n",
      ")\n",
      "> Training on CREATE TABLE sqlite_stat1(tbl,idx,stat)\n",
      "Adding ddl: CREATE TABLE sqlite_stat1(tbl,idx,stat)\n"
     ]
    }
   ],
   "source": [
    "pack = VannaPack(\n",
    "    openai_api_key=\"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\",\n",
    "    sql_db_url=\"chinook.db\",\n",
    "    openai_model=\"gpt-3.5-turbo\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e63e8a61-1fcf-40f9-b851-654cceb4d21e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Number of requested results 10 is greater than number of elements in index 9, updating n_results = 9\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using model gpt-3.5-turbo for 1401.5 tokens (approx)\n"
     ]
    }
   ],
   "source": [
    "response = pack.run(\"What are some sample albums?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c5815dbc-408a-49fc-b471-2863921a764f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Title\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8346e897-b8d3-4a76-965f-42d3836c052a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Number of requested results 10 is greater than number of elements in index 9, updating n_results = 9\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using model gpt-3.5-turbo for 1401.5 tokens (approx)\n"
     ]
    }
   ],
   "source": [
    "tmp = pack.get_modules()[\"vanna_query_engine\"].vn.ask(\n",
    "    \"What are some sample albums?\", visualize=False, print_results=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e7a87edb-b285-4f79-8cca-110a73758880",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'NoneType'>\n"
     ]
    }
   ],
   "source": [
    "print(type(tmp))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca8c0501-4ab7-4e23-9302-735fc6aa7b78",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_hub",
   "language": "python",
   "name": "llama_hub"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
